from pycocotools.coco import COCO
from .build import DatasetLoader
import os


@DatasetLoader.register()
class COCODataset:
    @classmethod
    def build(cls, dataset_config):
        dataset_root = dataset_config['dataset_root']
        dataset_json_file = dataset_config['dataset_json_file']

        coco = COCO(dataset_root, dataset_json_file)

        dataset = []

        for image_id in coco.getImgIds():
            image_info = coco.loadImgs(image_id)[0]
            image_path = os.path.join(dataset_root, image_info['file_name'])
            height = image_info['height']
            width = image_info['width']

            annotations_ids = coco.getAnnIds(imgIds=image_id, iscrowd=False)
            coco_annotations = coco.loadAnns(annotations_ids)

            objects_list = []
            for annotation in coco_annotations:
                bbox = annotation['bbox']
                bbox = [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]
                bbox[0] = max(0, bbox[0])
                bbox[1] = max(0, bbox[1])
                bbox[2] = min(width, bbox[2])
                bbox[3] = min(height, bbox[3])

                # some annotations have basically no width / height, skip them
                if bbox[2] - bbox[0] < 1 or bbox[3] - bbox[1] < 1:
                    continue

                label = coco.loadCats(annotation['category_id'])[0]['name']

                object_dict = {
                    'label': label,
                    'bounding_box': [float(i) for i in bbox],
                }
                objects_list.append(object_dict)

            # to uniform interface
            objects = []
            for object in objects_list:
                object_dict = {
                    'box': object['bounding_box'],
                    'category': object['label']
                }
                objects.append(object_dict)

            record = {
                'path': image_path,
                'objects': objects,
                'params': {
                    'height': height,
                    'width': width,
                },
                'transforms': [],
            }
            dataset.append(record)
        return dataset
